import argparse
import os
from pathlib import Path


project_dir = Path(__file__).resolve().parents[1]


def main():
    parser = argparse.ArgumentParser(description="Combine I/M/R Model Score.")
    parser.add_argument('--dataset', type=str, choices=['ucf101', 'hmdb51', 'ylimed', 'trecmed'],
                        help='dataset name.')
    parser.add_argument('--debug', action='store_true', help='Dont call system(), just show cmd.')
    # TODO: 保证参数组中至少传递了其中一个参数
    parser.add_argument('--iframe', type=str, default=None, help='iframe project name.')
    parser.add_argument('--mv', type=str, default=None, help='mv project name.')
    parser.add_argument('--res', type=str, default=None, help='residual project name.')
    args = parser.parse_args()
    class_num = {'ucf101': 101, 'hmdb51': 51, 'ylimed': 11, 'trecmed': 21}[args.dataset]

    # The simplest is the best
    combine = []  # 收集有意义的combine比例
    dirty = []
    collect = []
    if args.iframe:
        scores_dir = project_dir.joinpath('out', args.iframe, 'scores')
        npz_i = f'{scores_dir}/{args.dataset}_best_iframe_model__scores.npz'
        dirty.append(npz_i)
        collect.append(args.iframe)
        combine.append((1, 0, 0))
    if args.mv:
        scores_dir = project_dir.joinpath('out', args.mv, 'scores')
        npz_m = f'{scores_dir}/{args.dataset}_best_mv_model__scores.npz'
        dirty.append(npz_m)
        collect.append(args.mv)
        combine.append((0, 1, 0))
    if args.res:
        scores_dir = project_dir.joinpath('out', args.res, 'scores')
        npz_r = f'{scores_dir}/{args.dataset}_best_residual_model__scores.npz'
        dirty.append(npz_r)
        collect.append(args.res)
        combine.append((0, 0, 1))

    if args.iframe and args.mv:
        combine.append((1, 1, 0))
    if args.iframe and args.res:
        combine.append((1, 0, 1))
    if args.mv and args.res:
        combine.append((0, 1, 1))
    if args.iframe and args.mv and args.res:
        combine.append((1, 1, 1))
        combine.append((2, 1, 1))

    # None, [], {}, (), 0
    # 等一系列代表空和无的对象会被转换成False。除此之外的其它对象都会被转化成True.
    if not dirty:
        print('At least one --iframe/--mv/--res')
        return
    if args.iframe is None:
        npz_i = dirty[0]
    if args.mv is None:
        npz_m = dirty[0]
    if args.res is None:
        npz_r = dirty[0]

    # 当三个scores文件在相同位置, 就将结果放在此位置, 这样比较合理.
    if not (Path(npz_i).parent == Path(npz_m).parent == Path(npz_r).parent):
        scores_dir = project_dir.joinpath('out', 'combine', '-'.join(collect))
        scores_dir.mkdir(parents=True, exist_ok=True)

    # output = project_dir.joinpath('out', args.name)
    # weights_out = output.joinpath('weights')
    # scores_out = output.joinpath('scores')
    # log_out = output.joinpath('log')

    # TODO: 如果是多个文件夹的融合, 如何放置combine结果, 现在是以res/mv/i的优先级保持
    cmd = fr"""python {project_dir}/combine.py \
        --iframe    {npz_i} \
        --mv        {npz_m} \
        --res       {npz_r} \
        --num-class {class_num} \
        --combine "{combine}" \
        > {scores_dir}/combine.out 2>&1
    """
    print(cmd)
    if not args.debug:
        os.system(cmd)


if __name__ == "__main__":
    main()
